#include "sat.h"
#include <cstdio>
#include <cstdlib>

void setup_cuda(SAT_test_t *t);
void setup_buffers(SAT_test_t * t, float4(*gen)(int i));
void warmup(SAT_test_t * t, int num_iters);
bool run(SAT_test_t * t, int num_iters);
bool check(SAT_test_t *t);
	void SAT_gold(SAT_test_t *t);
void post(SAT_test_t *t, int run_num_iters, bool result);
void cleanup(SAT_test_t *t);

/////////////////////////////////////////////////////

void setup(SAT_test_t * t, float4(*gen)(int i))
{
	// malloc host memory

	for (int i = 0; i < t->num * t->num; i++) {
		t->hst_src[i] = gen(i);
	}
	setup_cuda(t);
	// malloc device memory for src and dst (i.e. result)

}

void setup_buffers(SAT_test_t * t, float4(*gen)(int i))
{
	// malloc host memory
	t->hst_src = (float4 *)malloc(t->num * t->num * sizeof(float4));
	t->hst_dst = (float4 *)malloc(t->num * t->num * sizeof(float4));
	for (int i = 0; i < t->num * t->num; i++) {
		t->hst_src[i] = gen(i);
	}

	// create GPU buffers
	printf("Creating GPU buffers...\n");
	cutilSafeCall(cudaMalloc(&(t->dev_src), t->num * t->num * sizeof(float4)));
	cutilSafeCall(cudaMalloc(&(t->dev_dst), t->num * t->num * sizeof(float4)));
}

void setup_cuda(SAT_test_t *t)
{
	// TODO
} 

float4 gen0(int i) {
	return make_float4(1.0f);
}

float4 gen1(int i) {
	return make_float4(1.0f*i);
}

float4 gen2(int i) {
	switch(i) {
		case 0 : return make_float4(2.0f); break;
		case 1 : return make_float4(1.0f); break;
		case 2 : return make_float4(0.0f); break;
		case 3 : return make_float4(0.0f); break;
		case 4 : return make_float4(0.0f); break;
		case 5 : return make_float4(1.0f); break;
		case 6 : return make_float4(2.0f); break;
		case 7 : return make_float4(0.0f); break;
		case 8 : return make_float4(1.0f); break;
		case 9 : return make_float4(2.0f); break;
		case 10 : return make_float4(1.0f); break;
		case 11 : return make_float4(0.0f); break;
		case 12 : return make_float4(1.0f); break;
		case 13 : return make_float4(1.0f); break;
		case 14 : return make_float4(0.0f); break;
		case 15 : return make_float4(2.0f); break;
		default : return make_float4(0.0f);
	}
}

float4 gen3(int i) {
	return make_float4(0.0f);
}

float4 gen4(int i) {
	return make_float4(cos(i/1000.0f));
}

float4 gen5(int i) {
	return make_float4(sin(i/100.0f));
}

void warmup(SAT_test_t * t, int num_iters)
{
	for (int i = 0; i < num_iters; i++) {
		launch_SAT(t);
		printf("Warmup %d/%d complete\n", i+1, num_iters);
	}
}

bool run(SAT_test_t * t, int num_iters)
{
	t->time = t->transpose_time = t->scan_time = 0.0;

	for (int i = 0; i < num_iters; i++) {
		double time = launch_SAT(t);
		t->time += time;
		if (!check(t)) 
			return false;
		printf("Run %d complete\n", i+1, num_iters);
	}
	return true;
}

bool check(SAT_test_t *t)
{
	float EPSILON = 2.0f;

	SAT_gold(t);
	float4 * temp = (float4 *) malloc (t->num * t->num * sizeof(float4));
	cutilSafeCall(cudaMemcpy(temp, t->dev_dst, t->num * t->num * sizeof(float4),
		cudaMemcpyDeviceToHost));
	for (int i = 0; i < t->num * t->num; i++) {
		float4 hst = t->hst_dst[i];
		float4 dev = temp[i];
		if (abs(hst.x - dev.x) >= EPSILON ||
			abs(hst.y - dev.y) >= EPSILON ||
			abs(hst.z - dev.z) >= EPSILON ||
			abs(hst.w - dev.w) >= EPSILON ) {
				fprintf(stderr, "check failed at %d: [%f, %f, %f, %f] != [%f, %f, %f, %f]\n", 
					i,
					hst.x, hst.y, hst.z, hst.w,
					dev.x, dev.y, dev.z, dev.w);
				free(temp);
				return false;
		}
	}
	free(temp);
	return true;
}

void SAT_gold(SAT_test_t * t)
{
	for (int y = 0; y < t->num; y++) {
		for (int x = 0; x < t->num; x++) {
			int i = y * t->num + x;
			float4 me = t->hst_src[i];
			float4 below = (y < t->num-1) ? t->hst_dst[(y+1)*t->num + x] : make_float4(0.0f);
			float4 above = (y > 0) ? t->hst_dst[(y-1)*t->num + x] : make_float4(0.0f);
			float4 left = (x > 0) ? t->hst_dst[y*t->num + x-1] : make_float4(0.0f);
			float4 below_left = (y < t->num-1 && x > 0) ? 
				t->hst_dst[(y+1)*t->num + x-1] : make_float4(0.0f);
			float4 above_left = (y > 0 && x > 0) ? 
				t->hst_dst[(y-1)*t->num + x-1] : make_float4(0.0f);

			t->hst_dst[i].x = me.x + above.x + left.x - above_left.x;
			t->hst_dst[i].y = me.y + above.y + left.y - above_left.y;
			t->hst_dst[i].z = me.z + above.z + left.z - above_left.z;
			t->hst_dst[i].w = me.w + above.w + left.w - above_left.w;
		}
	}
}

void print_arr(SAT_test_t * t, float4 *arr)
{
	for (int i = 0; i < t->num * t->num; i++) {
		if (i % t->num == 0)
			printf("[");
		printf("%3.0f", arr[i].x);
		if (i % t->num < t->num-1)
			printf(", ");
		else
			printf("]\n");
	}
}

void post(SAT_test_t * t, int run_num_iters, bool result)
{
	printf("----------------RESULTS----------------\n");
	printf("PASS? %s\n", result ? "YES" : "NO");
	printf("Size: %d x %d\n", t->num, t->num);
	printf("Time: %f\n", t->time / run_num_iters);
}

void output(SAT_test_t *t, FILE *csv, int run_num_iters) {
	fprintf(csv, "%u, %u, %Lf, %Lf, %Lf\n",
		t->num,
		t->blockSize,
		t->scan_time/(2 * run_num_iters),
		t->transpose_time/(2 * run_num_iters),
		t->time / run_num_iters);
	// / 2 since scan and transpose done twice
}

void cleanup(SAT_test_t * t)
{
	free(t->hst_src);
	free(t->hst_dst);
	cutilSafeCall(cudaFree(t->dev_src));
	cutilSafeCall(cudaFree(t->dev_dst));
}

void test_SAT(int sizeMin, int sizeMax,
			  int warmup_num_iters, int run_num_iters, 
			  int blockMin, int blockMax, 
			  bool split, char *filename, float4(*gen)(int i))
{
	SAT_test_t t;
	t.num = sizeMin;
	t.blockSize = blockMin;
	t.split = split;
	bool result = false;

	FILE * csv = fopen(filename, "a");
	if (csv == NULL)
		fprintf(stderr, "can't open output file\n");

	if (csv != NULL)
		fprintf(csv, "size, block size, scan time, transpose time, total time\n");

	setup_cuda(&t);
	setup_buffers(&t, gen);

	warmup(&t, warmup_num_iters);

	for (t.num = sizeMin; t.num <= sizeMax; t.num <<= 1) {
		for (t.blockSize = blockMin; t.blockSize <= blockMax; t.blockSize += 128) {
			setup_buffers(&t, gen);
			result = run(&t, run_num_iters);
			post(&t, run_num_iters, result);
			if (result && csv != NULL)
				output(&t, csv, run_num_iters);
			cleanup(&t);
		}
	}

	if (csv != NULL)
		fclose(csv);
}

int main(int argc, char *argv[])
{
	char *filename = "sat_cuda.csv";
	test_SAT(256, 4096, 1, 5, 128, 256, true, filename, gen4);
	test_SAT(256, 4096, 1, 5, 128, 256, false, filename, gen4);
	return 0;
}
